package cn.seecoder.ai.algorithm.ml.classification

import cn.seecoder.ai.algorithm.ml.BaseClassification
import cn.seecoder.ai.enums.TrainParamEnum
import cn.seecoder.ai.model.bo.TrainParamsBO
import cn.seecoder.ai.utils.TrainParamReader.readParam
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.classification.DecisionTreeClassifier

import java.util


// 决策树算法
class DecisionTreeClassification extends BaseClassification{

  override def buildMachineLearningStage(trainParams: TrainParamsBO): PipelineStage = {


    val dt = new DecisionTreeClassifier()
      .setLabelCol("label")
      .setFeaturesCol("indexedFeatures")
      .setSeed(trainParams.seed)
      .setMaxDepth(trainParams.maxDepth)
      .setMaxBins(trainParams.maxBins)
      .setMinInfoGain(trainParams.minInfoGain)


    dt
  }


}
